package com.baomidou.plugin.idea.mybatisx.advance.provider;

import com.baomidou.plugin.idea.mybatisx.advance.UserMapKeys;
import com.baomidou.plugin.idea.mybatisx.advance.processor.clazz.MPClassProcessor;
import com.baomidou.plugin.idea.mybatisx.advance.core.GenericServiceLocator;
import com.baomidou.plugin.idea.mybatisx.advance.processor.field.MPFieldProcessor;
import com.baomidou.plugin.idea.mybatisx.advance.util.PsiClassUtil;
import com.intellij.openapi.diagnostic.Logger;
import com.intellij.psi.*;
import com.intellij.psi.augment.PsiAugmentProvider;
import org.jetbrains.annotations.NotNull;

import java.util.*;

/**
 * Provides support for lombok generated elements
 *
 * @author Plushnikov Michail
 */
public class MPAugmentProvider extends PsiAugmentProvider {
  private static final Logger LOG = Logger.getInstance(MPAugmentProvider.class.getName());

  private final Collection<MPFieldProcessor> allFieldHandlers;
  private final Collection<MPClassProcessor> allClassHandlers;

  public MPAugmentProvider() {
    List<MPClassProcessor> lombokClassProcessors = GenericServiceLocator.locateAll(MPClassProcessor.class);
    List<MPFieldProcessor> lombokFieldProcessors = GenericServiceLocator.locateAll(MPFieldProcessor.class);

    allClassHandlers = new HashSet<MPClassProcessor>(lombokClassProcessors);
    allFieldHandlers = new HashSet<MPFieldProcessor>(lombokFieldProcessors);
  }

  @NotNull
  @Override
  public <Psi extends PsiElement> List<Psi> getAugments(@NotNull PsiElement element, @NotNull Class<Psi> type) {
    // Expecting that we are only augmenting an PsiClass
    // Don't filter !isPhysical elements or code autocompletion will not work
    if (!(element instanceof PsiClass) || !element.isValid()) {
      //noinspection unchecked
      return Collections.emptyList();
    }
    final List<Psi> result = new ArrayList<Psi>();
    final PsiClass psiClass = (PsiClass) element;

    if (type.isAssignableFrom(PsiField.class)) {
      LOG.info("collect field of class: " + psiClass.getQualifiedName());
      processPsiClassAnnotations(result, psiClass, type);

    } else if (type.isAssignableFrom(PsiMethod.class)) {
      LOG.info("collect methods of class: " + psiClass.getQualifiedName());

      cleanAttributeUsage(psiClass);
      processPsiClassAnnotations(result, psiClass, type);
      processPsiClassFieldAnnotation(result, psiClass, type);
    }
    return result;
  }

  protected void cleanAttributeUsage(PsiClass psiClass) {
    for (PsiField psiField : PsiClassUtil.collectClassFieldsIntern(psiClass)) {
      UserMapKeys.removeAllUsagesFrom(psiField);
    }
  }

  private <Psi extends PsiElement> void processPsiClassAnnotations(@NotNull List<Psi> result, @NotNull PsiClass psiClass, @NotNull Class<Psi> type) {
    final PsiModifierList modifierList = psiClass.getModifierList();
    if (modifierList != null) {
      for (PsiAnnotation psiAnnotation : modifierList.getAnnotations()) {
        processClassAnnotation(psiAnnotation, psiClass, result, type);
      }
    }
  }

  private <Psi extends PsiElement> void processClassAnnotation(@NotNull PsiAnnotation psiAnnotation, @NotNull PsiClass psiClass, @NotNull List<Psi> result, @NotNull Class<Psi> type) {
    for (MPClassProcessor classProcessor : allClassHandlers) {
      if (classProcessor.acceptAnnotation(psiAnnotation, type)) {
        classProcessor.process(psiClass, psiAnnotation, result);
      }
    }
  }

  protected <Psi extends PsiElement> void processPsiClassFieldAnnotation(@NotNull List<Psi> result, @NotNull PsiClass psiClass, @NotNull Class<Psi> type) {
    for (PsiField psiField : psiClass.getFields()) {
      processField(result, psiField, type);
    }
  }

  protected <Psi extends PsiElement> void processField(@NotNull List<Psi> result, @NotNull PsiField psiField, @NotNull Class<Psi> type) {
    final PsiModifierList modifierList = psiField.getModifierList();
    if (modifierList != null) {
      for (PsiAnnotation psiAnnotation : modifierList.getAnnotations()) {
        processFieldAnnotation(psiAnnotation, psiField, result, type);
      }
    }
  }

  private <Psi extends PsiElement> void processFieldAnnotation(@NotNull PsiAnnotation psiAnnotation, @NotNull PsiField psiField, @NotNull List<Psi> result, @NotNull Class<Psi> type) {
    for (MPFieldProcessor fieldProcessor : allFieldHandlers) {
      if (fieldProcessor.acceptAnnotation(psiAnnotation, type)) {
        fieldProcessor.process(psiField, psiAnnotation, result);
      }
    }
  }

}
